package interval

// SegmentTree 线段树结构
type SegmentTree struct {
	n    int
	tree []int64
	lazy []int64
}

// NewSegmentTree 创建线段树
func NewSegmentTree(n int) *SegmentTree {
	tree := &SegmentTree{
		n:    n,
		tree: make([]int64, 4*n),
		lazy: make([]int64, 4*n),
	}
	return tree
}

// Build 从数组构建线段树
func (st *SegmentTree) Build(arr []int64) {
	st.buildTree(arr, 1, 0, st.n-1)
}

func (st *SegmentTree) buildTree(arr []int64, node, start, end int) {
	if start == end {
		st.tree[node] = arr[start]
		return
	}

	mid := start + (end-start)/2
	leftNode := 2 * node
	rightNode := 2*node + 1

	st.buildTree(arr, leftNode, start, mid)
	st.buildTree(arr, rightNode, mid+1, end)

	st.tree[node] = st.tree[leftNode] + st.tree[rightNode]
}

// Update 区间更新
func (st *SegmentTree) Update(left, right int, val int64) {
	st.updateRange(1, 0, st.n-1, left, right, val)
}

func (st *SegmentTree) updateRange(node, start, end, left, right int, val int64) {
	if st.lazy[node] != 0 {
		st.tree[node] += st.lazy[node] * int64(end-start+1)
		if start != end {
			st.lazy[2*node] += st.lazy[node]
			st.lazy[2*node+1] += st.lazy[node]
		}
		st.lazy[node] = 0
	}

	if start > end || start > right || end < left {
		return
	}

	if start >= left && end <= right {
		st.tree[node] += val * int64(end-start+1)
		if start != end {
			st.lazy[2*node] += val
			st.lazy[2*node+1] += val
		}
		return
	}

	mid := start + (end-start)/2
	st.updateRange(2*node, start, mid, left, right, val)
	st.updateRange(2*node+1, mid+1, end, left, right, val)

	st.tree[node] = st.tree[2*node] + st.tree[2*node+1]
}

// Query 区间查询
func (st *SegmentTree) Query(left, right int) int64 {
	return st.queryRange(1, 0, st.n-1, left, right)
}

func (st *SegmentTree) queryRange(node, start, end, left, right int) int64 {
	if start > end || start > right || end < left {
		return 0
	}

	if st.lazy[node] != 0 {
		st.tree[node] += st.lazy[node] * int64(end-start+1)
		if start != end {
			st.lazy[2*node] += st.lazy[node]
			st.lazy[2*node+1] += st.lazy[node]
		}
		st.lazy[node] = 0
	}

	if start >= left && end <= right {
		return st.tree[node]
	}

	mid := start + (end-start)/2
	leftSum := st.queryRange(2*node, start, mid, left, right)
	rightSum := st.queryRange(2*node+1, mid+1, end, left, right)

	return leftSum + rightSum
}
